

#include "muxer_and_demuxer.hpp"
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <functional>
#include <unistd.h>
#include "define.h"
#include "scope_guard.h"
#include "simple_filter.h"
#include "frame_operate.h"

extern "C" {

    #include <libavformat/avformat.h>
    #include <libavformat/avio.h>
    #include <libavcodec/avcodec.h>
    #include <libavutil/frame.h>
    #include <libavcodec/codec.h>
    #include <libavcodec/packet.h>
    #include <libavutil/avutil.h>
    #include <libavcodec/codec_id.h>
    #include <libavcodec/codec_par.h>
    #include <libavutil/mathematics.h>
    #include <libavutil/rational.h>
    #include <libavutil/time.h>
    #include <libavfilter/buffersrc.h>
    #include <libavfilter/buffersink.h>
    #include <libavfilter/avfilter.h>
}

static bool CreateCodecContext(AVFormatContext * format_ctx, int type, OUT AVCodecContext ** codec_ctx) 
{

    int index = av_find_best_stream(format_ctx, AVMediaType(type), -1, -1, NULL, 0);
    if (index >= 0) 
    {

        AVStream * stream = format_ctx->streams[index];

        const AVCodec * codec = avcodec_find_decoder((AVCodecID)stream->codecpar->codec_id);
        *codec_ctx = avcodec_alloc_context3(codec);

        avcodec_parameters_to_context(*codec_ctx, stream->codecpar);

        int result = avcodec_open2(*codec_ctx, codec, NULL);
        if (result < 0) 
        {
            std::cout << "Error: avcodec open 2 " << std::endl;
        }
    }
    return index >= 0;
}

static bool CreateNewStream(AVFormatContext * input_format_ctx, int type, AVFormatContext * output_format_ctx)
{
    int index = av_find_best_stream(input_format_ctx, AVMediaType(type), -1, -1, NULL, 0);
    if (index >= 0)
    {
        AVStream * stream = input_format_ctx->streams[index];
        const AVStream * new_stream = avformat_new_stream(output_format_ctx, NULL);

        avcodec_parameters_copy(new_stream->codecpar, stream->codecpar);

        return true;
    }
    else 
    {
        return false;
    }
}

static bool CloseCodecContext(AVCodecContext * codec_context) 
{
    avcodec_close(codec_context);
    return true;
}

static bool CloseFormatContent(AVFormatContext * format_context) 
{
    avformat_close_input(&format_context);
    return true;
}

static bool CloseFrame(AVFrame * frame) 
{
    av_frame_free(&frame);
    return true;
}

static bool ClosePacket(AVPacket * packet) 
{
    av_packet_unref(packet);
    av_packet_free(&packet);
    return true;
}

static int ProcessVideoData(AVFrame * frame, void * user_data) 
{

    std::cout << "video data count----->video" << std::endl;
    hnhj::ProcessYUV420(frame, ((DecoderFile*)user_data)->m_file_handle);
    return 0;
}

static int ProcessAudioData(AVFrame * frame, void * user_data) 
{

    std::cout << "audio data count------>audio" << std::endl;
    WriteAudioData(frame, ((DecoderFile*)user_data)->m_file_handle);
    return 0;
}

static bool FilterFrame(AVFrame * frame, SimpleFilter * filter_data, CallbackData * callback)
{
    int result = av_buffersrc_add_frame_flags(filter_data->m_buffer_filter_ctx, frame, AV_BUFFERSRC_FLAG_KEEP_REF);
    if (result < 0)
    {
        std::cout << "Error: avbuffersrc add frame flags" << std::endl;
        return false;
    }
    while(true)
    {
        result = av_buffersink_get_frame(filter_data->m_sinker_filter_ctx, filter_data->m_output_frame);
        if (result < 0)
        {
            if (result == AVERROR_EOF || AVERROR(result) == EAGAIN) 
            {
                break;
            } 
            else 
            {
                std::cout << "Error: avcodec receive frame : " << av_err2str(result) << std::endl;
                return false;
            }
        }
        if (callback && callback->m_func)
        {
            callback->m_func(filter_data->m_output_frame, callback->m_user_data);
        }
        av_frame_unref(filter_data->m_output_frame);
    }
    return true;
}

static bool ProcessFilterFrame(AVPacket * packet, AVCodecContext * codec_context, AVFrame * frame, SimpleFilter * filter_data, CallbackData * callback)
{
    if (codec_context->codec_type == AVMEDIA_TYPE_AUDIO)
    {
        return true;
    }
    int result = avcodec_send_packet(codec_context, packet);
    if (result < 0) {

        std::cout << "Error: avcodec send packet : " << av_err2str(result) << std::endl;
        return false;
    }
    while (true) 
    {
        result = avcodec_receive_frame(codec_context, filter_data->m_input_frame);
        if (result < 0) 
        {
            if (result == AVERROR_EOF || AVERROR(result) == EAGAIN) 
            {
                break;
            } 
            else 
            {
                std::cout << "Error: avcodec receive frame : " << av_err2str(result) << std::endl;
                return false;
            }
        }
        if (codec_context->codec_type == AVMEDIA_TYPE_VIDEO)
        {
            return FilterFrame(filter_data->m_input_frame, filter_data, callback);
        } 
        else
        {
            // process frame data
            if (callback && callback->m_func)
            {
                callback->m_func(frame, callback->m_user_data);   
            }
        }
    }
    return true;
}

static bool ProcessFrame(AVPacket * packet, AVCodecContext * codec_context, AVFrame * frame, CallbackData * callback) 
{
    int result = avcodec_send_packet(codec_context, packet);
    if (result < 0) {

        std::cout << "Error: avcodec send packet : " << av_err2str(result) << std::endl;
        return false;
    }
    while (true) 
    {
        result = avcodec_receive_frame(codec_context, frame);
        if (result < 0) 
        {
            if (result == AVERROR_EOF || AVERROR(result) == EAGAIN) 
            {
                break;
            } 
            else 
            {
                std::cout << "Error: avcodec receive frame : " << av_err2str(result) << std::endl;
                return false;
            }
        }
        // process frame data
        if (callback && callback->m_func)
        {
            callback->m_func(frame, callback->m_user_data);   
        }
    }
    return true;
}

static void ProcessDemuxer(Demuxer * data, SimpleFilter * filter)
{
    AVFrame * frame = av_frame_alloc();
    AVPacket * packet = av_packet_alloc();
    while(av_read_frame(data->format_ctx, packet) >= 0) 
    {
        if (packet->stream_index == data->video_index) 
        {
            ProcessFilterFrame(packet, data->video_ctx, frame, filter, data->video_user_data);
        } 
        else if (packet->stream_index == data->audio_index) 
        {
            ProcessFilterFrame(packet, data->audio_ctx, frame, filter, data->audio_user_data);
        }
        av_packet_unref(packet);
    }
    ProcessFilterFrame(packet, data->audio_ctx, frame, filter, data->audio_user_data);  
    ProcessFilterFrame(packet, data->video_ctx, frame, filter, data->video_user_data);
}

void ProcessFile(const std::string &input_file) 
{

    AVFormatContext * format_ctx = nullptr;

    avformat_open_input(&format_ctx, input_file.c_str(), nullptr, nullptr);
    avformat_find_stream_info(format_ctx, nullptr);

    av_dump_format(format_ctx, 0, input_file.c_str(), 0);

    AVCodecContext * video_codec_ctx = nullptr;
    CreateCodecContext(format_ctx, AVMEDIA_TYPE_VIDEO, &video_codec_ctx);

    AVCodecContext * audio_codec_ctx = nullptr;
    CreateCodecContext(format_ctx, AVMEDIA_TYPE_AUDIO, &audio_codec_ctx);

    int video_index = 0;
    int audio_index = 0;
    for (int index = 0; index < format_ctx->nb_streams; index++) 
    {
        if (format_ctx->streams[index]->index == AVMEDIA_TYPE_VIDEO) 
        {
            video_index = index;
        } 
        else if (format_ctx->streams[index]->index == AVMEDIA_TYPE_AUDIO) 
        {
            audio_index = index;
        }
    }

    AVPacket * packet = av_packet_alloc();
    AVFrame  * frame  = av_frame_alloc();

    hnhj::ScopeGuard guard([&](){

        ClosePacket(packet);
        CloseFrame(frame);
        CloseCodecContext(video_codec_ctx);
        CloseCodecContext(audio_codec_ctx);
        CloseFormatContent(format_ctx);
    });

    // ready user data 
    DecoderFile audio_user_data = {"audio_test.pcm", NULL};
    DecoderFile video_user_data = {"video_test.yuv", NULL};

    audio_user_data.m_file_handle = fopen(audio_user_data.m_file_name.c_str(), "wb+");
    video_user_data.m_file_handle = fopen(video_user_data.m_file_name.c_str(), "wb+");

    Demuxer data;
    data.format_ctx  = format_ctx;
    data.video_ctx   = video_codec_ctx;
    data.video_index = video_index;
    data.audio_ctx   = audio_codec_ctx;
    data.audio_index = audio_index;

    data.audio_user_data = new CallbackData();
    data.audio_user_data->m_func      = std::bind(ProcessAudioData, std::placeholders::_1, std::placeholders::_2);
    data.audio_user_data->m_user_data = &audio_user_data;

    data.video_user_data = new CallbackData();
    data.video_user_data->m_func      = std::bind(ProcessVideoData, std::placeholders::_1, std::placeholders::_2);
    data.video_user_data->m_user_data = &video_user_data;

    SimpleFilter * filter = new SimpleFilter();

    // CreateFilter(filter, 1280, 720, "drawtext=fontsize=56:fontcolor=red:text='hello,world'");
    // CreateFilter(filter, 1280, 720, "zscale=640:360");

    // ProcessDemuxer(&data, filter);
}

void ProcessMuxer(const std::string &input_file)
{
    const std::string input_filename = input_file;
    const std::string output_filename = "./test.mp4";

    AVFormatContext * input_ctx  = NULL;
    AVFormatContext * output_ctx = NULL;

    int video_index = 0;
    int audio_index = 0;

    int result = 0;

    result = avformat_open_input(&input_ctx, input_filename.c_str(), NULL, NULL);
    if (result < 0)
    {
        std::cout << "Error: avformat open input" << std::endl;
        return;
    }
    avformat_find_stream_info(input_ctx, NULL);

    av_dump_format(input_ctx, 0, input_filename.c_str(), 0);

    avformat_alloc_output_context2(&output_ctx, NULL, "mp4", output_filename.c_str());

    int out_video_index = 0;
    int out_audio_index = 0;
    for (int index = 0; index < input_ctx->nb_streams; index++)
    {
        AVStream * stream = input_ctx->streams[index];
        if (stream->codecpar->codec_type == AVMEDIA_TYPE_AUDIO || stream->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
        {
            const AVCodec * codec = avcodec_find_encoder(stream->codecpar->codec_id);
            AVStream * new_stream = avformat_new_stream(output_ctx, codec);

            avcodec_parameters_copy(new_stream->codecpar, stream->codecpar);

            if (stream->codecpar->codec_type == AVMEDIA_TYPE_AUDIO)
            {
                audio_index = index;
            }
            if (stream->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
            {
                video_index = index;
            }
            if (new_stream->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
            {
                out_video_index = new_stream->index;
            }
            else
            {
                out_audio_index = new_stream->index;
            }
        }
    }
    if (!(output_ctx->oformat->flags & AVFMT_NOFILE))
    {
        avio_open(&output_ctx->pb, output_filename.c_str(), AVIO_FLAG_WRITE);
    }

    std::cout << "begin muxer" << std::endl;

    AVPacket * packet = av_packet_alloc();

    int video_frame_count = 0;
    int audio_frame_count = 0;

    avformat_write_header(output_ctx, NULL);

    int64_t start_time = av_gettime();

    while (av_read_frame(input_ctx, packet) >= 0) 
    {
        if (packet->pts == AV_NOPTS_VALUE)
        {
            int frame_count = 0;
            if (packet->stream_index == video_index)
            {
                frame_count = video_frame_count;
            }
            else
            {
                frame_count = audio_frame_count;
            }
            // frame_duration means a frame has how much time(ms)
            int64_t frame_duration = (double)AV_TIME_BASE/av_q2d(input_ctx->streams[packet->stream_index]->r_frame_rate);
            // time base translate
            packet->duration       = (double)frame_duration/av_q2d(input_ctx->streams[packet->stream_index]->time_base) * AV_TIME_BASE;
            //
            packet->pts            = (double)(frame_count * packet->duration)/(double)(av_q2d(input_ctx->streams[packet->stream_index]->time_base) * AV_TIME_BASE);
            // if not b frame dts equal pts
            packet->dts            = packet->pts;

        }
        if (output_ctx->oformat->flags & AVFMT_NOFILE)
        {
            AVRational time_base = input_ctx->streams[packet->stream_index]->time_base;
            AVRational time_base_p = {1, AV_TIME_BASE};
            int64_t pts_time = av_rescale_q(packet->pts, time_base, time_base_p);
            int64_t now_time = av_gettime() - start_time;
            if (pts_time > now_time)
            {
                av_usleep(pts_time - now_time);
            }
        }
        int out_stream_index = 0;
        if (packet->stream_index == video_index)
        {
            video_frame_count++;
            out_stream_index = out_video_index;
        }
        else if (packet->stream_index == audio_index)
        {
            audio_frame_count++;
            out_stream_index = out_audio_index;
        }
        AVStream * output_stream = output_ctx->streams[out_stream_index];
        AVStream * input_stream  = input_ctx->streams[packet->stream_index];

        // 按照输出format上下文转换时间基
        packet->pts = av_rescale_q_rnd(packet->pts, input_stream->time_base, output_stream->time_base, (AVRounding)(AV_ROUND_INF|AV_ROUND_PASS_MINMAX));
        packet->dts = av_rescale_q_rnd(packet->dts, input_stream->time_base, output_stream->time_base, (AVRounding)(AV_ROUND_INF|AV_ROUND_PASS_MINMAX));
        packet->duration = av_rescale_q(packet->duration, input_stream->time_base, output_stream->time_base);
        packet->pos = -1;

        av_interleaved_write_frame(output_ctx, packet);

        av_packet_unref(packet);
    }
    av_write_trailer(output_ctx);

    if (!(output_ctx->oformat->flags & AVFMT_NOFILE))
    {
        avio_closep(&output_ctx->pb);
    }

    av_packet_free(&packet);
    avformat_free_context(input_ctx);
    avformat_free_context(output_ctx);
}
